#!/usr/local/bin/python3
# -*- coding: utf-8 -*-

"""
@File    : bar_chart.py
@Author  : Link
@Time    : 2022/5/14 12:49
@Mark    : 
"""
import math

import numpy as np
import pandas as pd
from PySide2 import QtWidgets, QtCore
from typing import Union

from utils import CustomViewBox, PlotWidget, BarGraphItem


class BarChartPool:
    pass


class BarGraph(BarGraphItem):
    def mouseClickEvent(self, event):
        print("clicked")


class BarChartParams:
    start = 0
    stop = 0

    data_max = 10
    data_min = 0

    temp_df = None  # type:pd.DataFrame


class BarChart(QtCore.QObject):
    vb = None  # type:CustomViewBox
    pw = None  # type:PlotWidget
    bg1 = None  # type:BarGraph
    bg2 = None  # type:BarGraph  # 选取的

    ticks = None  # type:list
    df = None  # type:pd.DataFrame
    temp_df = None  # type:pd.DataFrame
    key = None  # type:str

    bg1_params = None  # type: BarChartParams
    bg2_params = None  # type: Union[BarChartParams, None]

    bar_width = 0  # 这个是柱状图的当前宽度
    bins_num = 60

    signal = QtCore.Signal(object)  # 用来将BarChart中操作过的数据进行一个传出

    def __init__(self, title: str = "PlotItem with CustomAxisItem, CustomViewBox"):
        super(BarChart, self).__init__()
        self.vb = CustomViewBox()
        self.pw = PlotWidget(viewBox=self.vb, enableMenu=False, title=title)

        self.bg1 = BarGraph(x0=[], y=[], y0=[], y1=[], width=[], brush='r')
        self.pw.addItem(self.bg1)
        self.bg2 = BarGraph(x0=[], y=[], y0=[], y1=[], width=[], brush=(217, 83, 25, 255))
        self.pw.addItem(self.bg2)

        self.bottom_axis = self.pw.getAxis("bottom")
        self.bottom_axis.setHeight(60)
        self.left_axis = self.pw.getAxis("left")
        self.left_axis.setWidth(60)
        self.bg1_params = BarChartParams()

    def init_signal(self):
        self.pw.sigRangeChanged.connect(self.set_range_data_to_chart)
        self.vb.select_signal.connect(self.select_range)

    def widget(self):
        return self.pw

    def set_bins(self, n):
        self.bins_num = n

    def set_data(self, key: str, df: pd.DataFrame, data_min=None, data_max=None):
        """

        :param key:
        :param df: 应该设置成 dict
        :param data_min:
        :param data_max:
        :return:
        """
        self.bg1_params.data_min, self.bg1_params.data_max, self.df, self.key = data_min, data_max, df, key
        if data_min is None:
            self.bg1_params.data_min = df[key].min()
        if data_max is None:
            self.bg1_params.data_max = df[key].max()

    def select_range(self, ax: QtCore.QRectF):
        """
        要使用当前的柱状图
        :param ax:
        :return:
        """
        if ax is None:
            self.bg2.setOpts(x0=[], y=[], y0=[], y1=[], width=[], brush=(217, 83, 25, 255))
            self.set_data_to_chart()
            return
        if self.df is None:
            self.bg2_params = None
            return
        start = math.ceil(ax.left() / self.bar_width)
        stop = math.ceil(ax.right() / self.bar_width)
        if start > len(self.ticks) or stop < 0:
            self.bg2_params = None
            return
        data_min = ax.top()
        data_max = ax.bottom()
        temp_df = self.df[(self.df[self.key] > data_min) & (self.df[self.key] < data_max)]
        if len(temp_df) == 0:
            self.bg2_params = None
            return
        # self.signal.emit("A")
        if self.bg2 is None:
            self.bg2 = BarGraph(x0=[], y=[], y0=[], y1=[], width=[], brush=(217, 83, 25, 255))
            self.pw.addItem(self.bg2)
        else:
            self.bg2.setOpts(x0=[], y=[], y0=[], y1=[], width=[], brush=(217, 83, 25, 255))
        self.bg2_params = BarChartParams()
        self.bg2_params.start, self.bg2_params.stop, self.bg2_params.data_max, self.bg2_params.data_min, \
        self.bg2_params.temp_df = start, stop, data_max, data_min, temp_df
        self.set_range_date_to_chart()
        # self.set_data_to_chart((217, 83, 25, 200))

    def set_range_date_to_chart(self, brush=(217, 83, 25, 255)):
        x0, y0, y1, y, width = [], [], [], [], []

        df_dict = {}
        for key, e_df in self.bg2_params.temp_df.groupby("GROUP"):
            if len(e_df) == 0:
                continue
            df_dict[key] = e_df
        list_bins = np.linspace(self.bg1_params.data_min, self.bg1_params.data_max, self.bins_num)
        for i in range(self.bg2_params.start - 1, self.bg2_params.stop):
            if i < 0:
                continue
            if i == len(self.ticks):
                break
            key = self.ticks[i]
            if key not in df_dict:
                continue
            temp_dis = df_dict[key][self.key].value_counts(bins=list_bins)
            for index, value in temp_dis.items():
                x0.append(i + 0.2)
                y0.append(index.left)
                y1.append(index.right)
                y.append((index.left + index.right) / 2)
                width.append(value)
        x0 = np.array(x0) * self.bar_width
        self.bg2.setOpts(x0=x0, y=y, y0=y0, y1=y1, width=width, brush=brush)

    def set_range_data_to_chart(self, a, ax):
        if self.df is None:
            return
        if ax[1][1] - ax[1][0] == self.bg1_params.data_max - self.bg1_params.data_min:
            """ 性能会好些,体验会差点 """
            return
        self.bg1_params.data_min, self.bg1_params.data_max = ax[1]
        if self.bg2_params is None:
            # 如果有range就先放range的
            self.set_data_to_chart()
        else:
            self.set_data_to_chart()
            self.set_range_date_to_chart()

    def set_data_to_chart(self, brush=(217, 83, 25, 140)):
        """
        设定和更新界面
        :param brush:
        :return:
        """
        temp_df = self.df[
            (self.df[self.key] > self.bg1_params.data_min) & (self.df[self.key] < self.bg1_params.data_max)]
        if len(temp_df) == 0:
            return
        list_bins = np.linspace(self.bg1_params.data_min, self.bg1_params.data_max, self.bins_num)
        list_bins = list_bins.round(8)
        columns, c, x0, y0, y1, y, self.bar_width, width = [], 0, [], [], [], [], 0, []
        for key, e_df in temp_df.groupby("GROUP"):
            temp_dis = e_df[self.key].value_counts(bins=list_bins)
            columns.append(key)
            self.bar_width = max(temp_dis) if max(temp_dis) > self.bar_width else self.bar_width
            for index, value in temp_dis.items():
                x0.append(c + 0.2)
                y0.append(index.left)
                y1.append(index.right)
                y.append((index.left + index.right) / 2)
                width.append(value)
            c = c + 1
        x0 = np.array(x0) * self.bar_width
        self.ticks = columns
        ticks = [((idx + 0.2) * self.bar_width, label.replace("-", "\r\n")) for idx, label in enumerate(columns)]
        self.bg1.setOpts(x0=x0, y=y, y0=y0, y1=y1, width=width, brush=brush)
        self.bottom_axis.setTicks((ticks, []))
        self.vb.setXRange(0, c * self.bar_width * 1.2)


if __name__ == '__main__':
    """
    测试
    """
    import pyqtgraph as pg
    from pyqtgraph.dockarea import DockArea, Dock

    df = pd.read_csv("te/distribution_trans.csv")
    app = pg.mkQApp()
    win = QtWidgets.QMainWindow()
    area = DockArea()
    win.setCentralWidget(area)
    win.resize(1000, 500)

    bar_class = BarChart()
    # bar_class.signal.connect(lambda:print("A"))

    d1 = Dock("Dock1", size=(1, 1))
    area.addDock(d1, 'left')
    d1.addWidget(bar_class.widget())
    bar_class.set_data("2010:VOL_VREG1_LOW   JUDGE_VARIABLE_MS", df=df)
    bar_class.set_data_to_chart()
    win.show()

    bar_class.init_signal()
    pg.exec()
